package com.nbsaas.boot;

import javax.crypto.Mac;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.spec.KeySpec;
import java.util.Arrays;

/**
 * NovelFeistelNumHandle
 *
 * - 实现了用户给定的接口 NumHandle
 * - encode(Integer num) -> 32 hex chars (128-bit)
 * - decode(String keyHex) -> Integer (可逆)
 *
 * 设计要点：
 * - 使用 Feistel 网络 (ROUNDS 轮)，块大小 128-bit，左右各 64-bit
 * - 每轮的 round-function: HMAC-SHA256( roundKey, rightHalf || roundIndex ) -> 64-bit (截断)
 * - roundKey 由 PBKDF2(secret, salt || roundIndex) 或主键 + HMAC 派生，示例使用一次 PBKDF2 得到主 key，再用 HMAC-SHA256 为每轮派生
 * - 明文组织：16 bytes block; 我们把 Integer 放在 block 的低 4 字节（最后4字节），其余填固定常量（版本/tweak 可扩展）
 *
 * 警告：这是自定义构造，未经过形式证明。若用于安全关键系统请先审计。
 */
public class NovelFeistelNumHandle implements NumHandle {

    private static final int ROUNDS = 16; // 轮数（可调整），越多越安全但越慢
    private static final int KDF_ITER = 100_000; // PBKDF2 迭代次数（可调）
    private static final int KEY_BYTES = 32; // PBKDF2 输出字节 -> 256-bit 主键
    private static final byte[] KDF_SALT = "NovelFeistelStaticSalt_v1".getBytes(StandardCharsets.UTF_8);

    private final byte[] masterKey; // 256-bit 主键，从 secret 派生

    // 构造器：传入 secret（必须）
    public NovelFeistelNumHandle(String secret) {
        if (secret == null || secret.isEmpty()) throw new IllegalArgumentException("secret required");
        this.masterKey = kdfPbkdf2(secret.toCharArray(), KDF_SALT, KDF_ITER, KEY_BYTES);
    }

    // ========== NumHandle 接口实现 ==========

    @Override
    public String encode(Integer num) {
        if (num == null) throw new IllegalArgumentException("num required");
        // 构造 16 字节明文块
        byte[] plain = makePlainBlock(num);

        // 分左右 64-bit
        long L = bytesToLong(Arrays.copyOfRange(plain, 0, 8));
        long R = bytesToLong(Arrays.copyOfRange(plain, 8, 16));

        // Feistel 轮
        for (int round = 0; round < ROUNDS; round++) {
            long roundKey = deriveRoundValue(round, masterKey); // 64-bit value from HMAC
            long f = roundFunction(R, roundKey); // 64-bit
            long newL = R;
            long newR = L ^ f;
            L = newL;
            R = newR;
        }

        // 合并为 16 字节密文
        byte[] cipher = new byte[16];
        System.arraycopy(longToBytes(L), 0, cipher, 0, 8);
        System.arraycopy(longToBytes(R), 0, cipher, 8, 8);

        // 转 32 字符 hex（小写）
        return toHexLower(cipher);
    }

    @Override
    public Integer decode(String key) {
        if (key == null) throw new IllegalArgumentException("key required");
        if (key.length() != 32) throw new IllegalArgumentException("key must be 32 hex chars");

        byte[] cipher = fromHex(key);

        long L = bytesToLong(Arrays.copyOfRange(cipher, 0, 8));
        long R = bytesToLong(Arrays.copyOfRange(cipher, 8, 16));

        // 反向 Feistel
        for (int round = ROUNDS - 1; round >= 0; round--) {
            long roundKey = deriveRoundValue(round, masterKey);
            long f = roundFunction(L, roundKey); // note: symmetric usage
            long newR = L;
            long newL = R ^ f;
            L = newL;
            R = newR;
        }

        // 合并得到明文
        byte[] plain = new byte[16];
        System.arraycopy(longToBytes(L), 0, plain, 0, 8);
        System.arraycopy(longToBytes(R), 0, plain, 8, 8);

        // 从明文中提取 Integer（最后4字节）
        int recovered = ByteBuffer.wrap(Arrays.copyOfRange(plain, 12, 16)).getInt();
        return recovered;
    }

    // ========== 辅助函数 ==========

    // 明文格式：16 bytes
    // [0..7]  = fixed prefix (can be version/tweak)
    // [8..11] = zeros (可留作扩展)
    // [12..15] = 4 bytes of integer (big-endian)
    private static byte[] makePlainBlock(int num) {
        byte[] block = new byte[16];
        // prefix 用 SHA-256(masterKey||"prefix") 的前8字节这样更难预测（可选）
        // 为简单起见，这里使用固定常量前缀（可替换成基于 masterKey 的动态前缀）
        byte[] prefix = "NFH_v1__".getBytes(StandardCharsets.US_ASCII); // 8 bytes
        System.arraycopy(prefix, 0, block, 0, Math.min(8, prefix.length));
        // block[8..11] 留作0
        byte[] intBytes = ByteBuffer.allocate(4).putInt(num).array();
        System.arraycopy(intBytes, 0, block, 12, 4);
        return block;
    }

    // 轮函数 f(right, roundKey) -> 64-bit
    // 我们使用：HMAC-SHA256(roundKeyBytes, rightHalfBytes || roundIndex) -> 32 bytes -> 8 bytes truncated -> long
    private static long roundFunction(long right, long roundKey) {
        try {
            byte[] rightBytes = longToBytes(right);
            byte[] roundKeyBytes = longToBytes(roundKey);
            // HMAC-SHA256
            javax.crypto.Mac mac = javax.crypto.Mac.getInstance("HmacSHA256");
            SecretKeySpec kSpec = new SecretKeySpec(roundKeyBytes, "HmacSHA256");
            mac.init(kSpec);
            mac.update(rightBytes);
            byte[] h = mac.doFinal(); // 32 bytes
            // 取前8字节作为 64-bit 返回
            return bytesToLong(Arrays.copyOf(h, 8));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // 派生每一轮的 64-bit round value（基于 masterKey + round）
    private static long deriveRoundValue(int round, byte[] masterKey) {
        try {
            // 使用 HMAC-SHA256(masterKey, "ROUND" || round)
            Mac mac = Mac.getInstance("HmacSHA256");
            SecretKeySpec k = new SecretKeySpec(masterKey, "HmacSHA256");
            mac.init(k);
            mac.update("FEISTEL_ROUND".getBytes(StandardCharsets.UTF_8));
            mac.update((byte) round);
            byte[] out = mac.doFinal();
            // 取前8字节
            return bytesToLong(Arrays.copyOf(out, 8));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // PBKDF2 基于 secret 生成主键（masterKey）
    private static byte[] kdfPbkdf2(char[] password, byte[] salt, int iter, int keyLenBytes) {
        try {
            SecretKeyFactory skf = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
            KeySpec spec = new PBEKeySpec(password, salt, iter, keyLenBytes * 8);
            byte[] key = skf.generateSecret(spec).getEncoded();
            return key;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // 辅助：long <-> bytes（big-endian）
    private static byte[] longToBytes(long v) {
        ByteBuffer bb = ByteBuffer.allocate(8);
        bb.putLong(v);
        return bb.array();
    }
    private static long bytesToLong(byte[] b) {
        ByteBuffer bb = ByteBuffer.wrap(b);
        return bb.getLong();
    }

    // hex 编解码（小写）
    private static String toHexLower(byte[] b) {
        StringBuilder sb = new StringBuilder(b.length * 2);
        for (byte x : b) {
            sb.append(String.format("%02x", x));
        }
        return sb.toString();
    }
    private static byte[] fromHex(String s) {
        int len = s.length();
        byte[] out = new byte[len/2];
        for (int i = 0; i < len; i += 2) {
            out[i/2] = (byte)((Character.digit(s.charAt(i),16) << 4) + Character.digit(s.charAt(i+1),16));
        }
        return out;
    }
}
